# encoding: utf-8
"""
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    created by lei.chang on '18/12/2024'
    comment: 扣子平台相关
"""
import asyncio
import aiohttp
import ujson
import time
import uuid
from authlib.jose import jwt
from pydantic import Field
from laner.pydantic import BaseModel
from typing import Union

from config import Config
from project.core.function import is_url_urllib
from project.core.exception import BizError
from project.lib.common import logger
from project.core.base import UEnum, Element


class ObjectString(BaseModel):
    """
    """
    type: str = Field('text', title='多模态消息内容类型', description='text(文本类型)/file(文件类型)/image(图片类型)')
    text: str = Field('', title='文本内容')
    file_id: str = Field('', title='文件或图片内容的 ID')
    file_url: str = Field('', title='文件或图片内容的在线地址。必须是可公共访问的有效地址。', description='在 type 为 file 或 image 时，file_id 和 file_url 应至少指定一个。')


class EnterMessage(BaseModel):
    """
    """
    role: str = Field(..., title='发送消息的实体', description='user(用户)/assistant(智能体)')
    type: str = Field('question', title='消息类型', description='question/answer/function_call/tool_output/tool_response/follow_up/verbose')
    content: Union[str, list[ObjectString]] = Field('', title='消息的内容', description='支持纯文本、多模态（文本、图片、文件混合输入）、卡片等多种类型的内容')
    content_type: str = Field('text', title='消息内容的类型', description='text(文本)/object_string(多模态内容)')
    meta_data: dict = Field({}, title='创建消息时的附加消息，获取消息时也会返回此附加消息。')

    def dict(self, *args, **kwargs) -> dict:
        """
        :param args:
        :param kwargs:
        :return:
        """
        if isinstance(self.content, list):
            self.content = ujson.dumps(self.content, ensure_ascii=False)

        result = super(BaseModel, self).dict(*args, **kwargs)

        return result


class ChatEvent(UEnum):
    """ 流式事件
    """
    created = Element('conversation.chat.created', '创建对话', description='创建对话的事件，表示对话开始。')
    in_progress = Element('conversation.chat.in_progress', '处理状态', description='服务端正在处理对话。')
    message = Element('conversation.message.delta', '增量消息', description='通常是 type=answer 时的增量消息。')
    message_completed = Element('conversation.message.completed', 'message 已回复完成', description='通常是 type=answer 时的增量消息。')
    chat_completed = Element('conversation.chat.completed', 'message 对话完成')
    chat_failed = Element('conversation.chat.failed', '此事件用于标识对话失败。')
    requires_action = Element('conversation.chat.requires_action', '会话中断')
    error = Element('error', '会话异常')
    done = Element('done', '本次会话的流式返回正常结束。')


def get_enter_message(user_message='', files=None, images=None, meta_data=None):
    """
    :param user_message:
    :param files:
    :param images:
    :param meta_data:
    :return:
    """
    enter_message = EnterMessage(role='user')
    enter_message.meta_data = meta_data
    text_object_string = None
    if user_message:
        enter_message.content = user_message
        enter_message.content_type = 'text'

        text_object_string = ObjectString()
        text_object_string.type = 'text'
        text_object_string.text = user_message

    object_strings = list()
    # 文件列表
    if files is not None:
        for file in files:
            object_string = ObjectString(type='file')
            if is_url_urllib(file):
                object_string.file_url = file
            else:
                object_string.file_id = file

            object_strings.append(object_string)

    # 图片列表
    if images is not None:
        for image in images:
            object_string = ObjectString(type='image')
            if is_url_urllib(image):
                object_string.file_url = image
            else:
                object_string.file_id = image

            object_strings.append(object_string)

    if object_strings:
        enter_message.content_type = 'object_string'
        if text_object_string:
            object_strings.insert(0, text_object_string)
        enter_message.content = object_strings

    return enter_message.dict() if enter_message.content or enter_message.meta_data else {}


class Coze:

    def __init__(self, personal_token=None, public_key:str=None, private_key:str=None, oauth_app_id:str=None):
        """
        """
        self.personal_token = personal_token
        self.content_type = 'application/json'
        self.public_key = public_key
        self.private_key = private_key
        self.oauth_app_id = oauth_app_id

    async def _get_headers(self):
        """
        :return:
        """
        headers = {'Content-Type': 'application/json'}
        if self.public_key is not None and self.private_key is not None and self.oauth_app_id is not None:
            oauth_token = await self.__class__.get_oauth_token(public_key=self.public_key, private_key=self.private_key, oauth_app_id=self.oauth_app_id)
            headers['Authorization'] = f'Bearer {oauth_token}'
        elif self.personal_token is not None:
            # print("coze令牌使用个人临时令牌")
            headers['Authorization'] = f'Bearer {self.personal_token}'
        else:
            oauth_token = await self.__class__.get_oauth_token()
            headers['Authorization'] = f'Bearer {oauth_token}'

        return headers

    @classmethod
    async def get_oauth_token(cls, public_key=None, private_key=None, oauth_app_id=None):
        """ 取得jwt oauth token
        :return:
        """
        use_team_key = public_key is None or private_key is None or oauth_app_id is None
        if use_team_key:
            # print("coze令牌使用团队密钥获取")
            team = Config.COZE_OPTIONS['team']
            with open(team['private_key_file'], 'rb') as file:
                private_key_pem = file.read()

            private_key = private_key_pem.decode('utf-8')
            public_key = team['public_key']
            oauth_app_id = team['app_id']
        # else:
        #     print(f"coze令牌使用个人密钥获取。oauth_app_id: {oauth_app_id}")

        header = {
            'alg': 'RS256',
            'typ': 'JWT',
            'kid': public_key
        }

        now_timestamp = int(time.time())
        pay_load = {
            'iss': oauth_app_id,
            'aud': 'api.coze.cn',
            'iat': now_timestamp,
            'exp': now_timestamp + 86399,
            'jti': str(uuid.uuid4()).replace('-', '')
        }

        # 使用私钥生成签名
        try:
            jwt_token = jwt.encode(header, pay_load, private_key).decode('utf-8')
        except Exception as e:
            logger.error(f"jwt编码失败，ValueError: {str(e)}")
            raise BizError(message=f'生成jwt token失败，请检查密钥格式是否正确。')

        async with aiohttp.ClientSession() as session:
            async with session.post('https://api.coze.cn/api/permission/oauth2/token',
                                    headers={'Authorization': f'Bearer {jwt_token}'},
                                    json={
                                        'duration_seconds': 86399,
                                        'grant_type': 'urn:ietf:params:oauth:grant-type:jwt-bearer'
                                    }) as resp:
                if resp.status != 200:
                    error = f'扣子获取令牌失败 status: {resp.status} content: {(await resp.content.read())}'
                    if not use_team_key:
                        error = f'{error} 请检查公钥、私钥、app_id是否正确。'
                        # return await cls.get_oauth_token()
                    logger.error(message=error)
                    raise BizError(message=error)

                result = await resp.json()

        return result.get('access_token')

    async def create_conversation(self, user_message='', files=None, images=None, meta_data=None):
        """ 创建会话
        :param user_message:
        :param files:
        :param images:
        :param meta_data: 附加消息
        :return:
        """
        url = 'https://api.coze.cn/v1/conversation/create'

        headers = await self._get_headers()

        message = get_enter_message(user_message=user_message, files=files, images=images, meta_data=meta_data)
        params = {
            'messages': [message]
        }
        async with aiohttp.ClientSession() as session:
            async with session.post(url, headers=headers, json=params) as resp:
                if resp.status != 200:
                    error = f'扣子创建会话失败 url: {url} status: {resp.status} content: {(await resp.content.read())}'
                    logger.error(message=error)
                    raise BizError(message=error)

                result = await resp.json()

                if result.get('code', 0) != 0:
                    error = f'扣子创建会话失败 url: {url} code: {result.get("code", 0)} msg: {result.get("msg", "")}'
                    logger.error(message=error)
                    raise BizError(message=error)

        return result['data']

    async def retrieve_conversation(self, conversation_id):
        """ 查看会话
        :param conversation_id:
        :return:
        """
        url = 'https://api.coze.cn/v1/conversation/retrieve'

        headers = await self._get_headers()
        params = {
            'conversation_id': conversation_id
        }
        async with aiohttp.ClientSession() as session:
            async with session.get(url, headers=headers, params=params) as resp:
                if resp.status != 200:
                    error = f'扣子查看会话失败 url: {url} status: {resp.status} content: {(await resp.content.read())}'
                    logger.error(message=error)
                    raise BizError(message=error)

                result = await resp.json()

                if result.get('code', 0) != 0:
                    error = f'扣子查看会话失败 url: {url} code: {result.get("code", 0)} msg: {result.get("msg", "")}'
                    logger.error(message=error)
                    raise BizError(message=error)

        return result['data']

    async def chat(self,
                   bot_id,
                   user_id,
                   user_message,
                   conversation_id=None,
                   files=None,
                   images=True,
                   meta_data=None):
        """ 发起对话
        :return:
        """
        url = 'https://api.coze.cn/v3/chat'

        headers = await self._get_headers()
        if conversation_id:
            url = f'{url}?conversation_id={conversation_id}'

        body = dict()
        body['bot_id'] = bot_id
        body['user_id'] = user_id
        enter_message = get_enter_message(user_message, files, images, meta_data)
        body['additional_messages'] = [enter_message]
        body['stream'] = False
        body['auto_save_history'] = True
        if meta_data is not None:
            body['meta_data'] = meta_data

        async with aiohttp.ClientSession() as session:
            async with session.post(url, headers=headers, json=body) as resp:
                if resp.status != 200:
                    error = f'扣子发起会话失败 url: {url} status: {resp.status} content: {(await resp.content.read())}'
                    logger.error(message=error)
                    raise BizError(message=error)

                result = await resp.json()

                if result.get('code', 0) != 0:
                    error = f'扣子发起会话失败 url: {url} code: {result.get("code", 0)} msg: {result.get("msg", "")}'
                    logger.error(message=error)
                    raise BizError(message=error)

        return result['data']

    async def chat_streaming(self,
                             bot_id,
                             user_id,
                             user_message,
                             conversation_id=None,
                             files=None,
                             images=True,
                             meta_data=None,
                             custom_variables=None):
        """ 发起对话(流式)
        :return:
        """
        url = 'https://api.coze.cn/v3/chat'

        try:
            headers = await self._get_headers()
        except Exception as e:
            yield str(e)
            return

        if conversation_id:
            url = f'{url}?conversation_id={conversation_id}'

        body = dict()
        body['bot_id'] = bot_id
        body['user_id'] = user_id
        enter_message = get_enter_message(user_message, files, images, meta_data)
        body['additional_messages'] = [enter_message]
        body['stream'] = True
        body['auto_save_history'] = True
        if meta_data is not None:
            body['meta_data'] = meta_data
        body['custom_variables'] = custom_variables

        async with aiohttp.ClientSession() as session:
            async with session.post(url, headers=headers, json=body) as resp:
                if resp.status != 200:
                    error = f'扣子发起会话失败 url: {url} status: {resp.status} content: {(await resp.content.read())}'
                    logger.error(message=error)
                    raise BizError(message=error)

                try:
                    async for data in resp.content.iter_any():
                        data = data.decode('utf-8')
                        try:
                            result = ujson.loads(data)
                            if result.get('code') in (4101, 4100):
                                yield f"您的token没有权限访问Bot ID为{bot_id}的聊天。请验证资源ID或token是否正确。"
                                return
                            elif result.get('code') == 4015:
                                yield f"Bot ID为{bot_id}的机器人还没有发布到API通道。"
                                return
                        except:
                            pass
                        read_next_line = False
                        for line in data.split('\n'):
                            if not line:
                                continue
                            if line.startswith('event:') and str(ChatEvent.message.value) in line:
                                read_next_line = True

                            if line.startswith('data:') and read_next_line:
                                data = ujson.loads(line.replace('data:', ''))
                                read_next_line = False
                                yield data.get('content', '')
                except Exception as ex:
                    logger.critical(str(ex))
                    raise BizError(message=str(ex))

    async def retrieve_chat(self, conversation_id, chat_id):
        """
        :param conversation_id:
        :param chat_id:
        :return:
        """
        url = 'https://api.coze.cn/v3/chat/retrieve'
        headers = await self._get_headers()
        params = {
            'conversation_id': conversation_id,
            'chat_id': chat_id
        }

        async with aiohttp.ClientSession() as session:
            async with session.get(url, headers=headers, params=params) as resp:
                if resp.status != 200:
                    error = f'扣子查看会话详情失败 url: {url} status: {resp.status} content: {(await resp.content.read())}'
                    logger.error(message=error)
                    raise BizError(message=error)

                result = await resp.json()

                if result.get('code', 0) != 0:
                    error = f'扣子查看会话详情失败 url: {url} code: {result.get("code", 0)} msg: {result.get("msg", "")}'
                    logger.error(message=error)
                    raise BizError(message=error)

        return result['data']

    async def get_chat_message_list(self, conversation_id, chat_id):
        """
        :param conversation_id:
        :param chat_id:
        :return:
        """
        url = 'https://api.coze.cn/v3/chat/message/list'
        headers = await self._get_headers()
        params = {
            'conversation_id': conversation_id,
            'chat_id': chat_id
        }

        async with aiohttp.ClientSession() as session:
            async with session.get(url, headers=headers, params=params) as resp:
                if resp.status != 200:
                    error = f'扣子查看会话详情失败 url: {url} status: {resp.status} content: {(await resp.content.read())}'
                    logger.error(message=error)
                    raise BizError(message=error)

                result = await resp.json()

                if result.get('code', 0) != 0:
                    error = f'扣子查看会话详情失败 url: {url} code: {result.get("code", 0)} msg: {result.get("msg", "")}'
                    logger.error(message=error)
                    raise BizError(message=error)

        return result['data']

    async def run_workflow(self, workflow_id, parameters=None, bot_id=None, app_id=None, ext=None, is_async=False):
        """ 执行工作流
        :param workflow_id:
        :param parameters:
        :param bot_id:
        :param app_id:
        :param ext:
        :param is_async:
        :return:
        """
        url = 'https://api.coze.cn/v1/workflow/run'

        headers = await self._get_headers()
        params = {'workflow_id': workflow_id}
        if parameters is not None:
            params['parameters'] = parameters
        if bot_id is not None:
            params['bot_id'] = bot_id
        if app_id is not None:
            params['app_id'] = app_id
        if ext is not None:
            params['ext'] = ext
        params['is_async'] = is_async

        async with aiohttp.ClientSession() as session:
            async with session.post(url, headers=headers, json=params) as resp:
                if resp.status != 200:
                    error = f'扣子执行工作流失败 url: {url} status: {resp.status} content: {(await resp.content.read())}'
                    logger.error(message=error)
                    raise BizError(message=error)

                result = await resp.json()

                if result.get('code', 0) != 0:
                    error = f'扣子执行工作流失败 url: {url} code: {result.get("code", 0)} msg: {result.get("msg", "")}'
                    logger.error(message=error)
                    raise BizError(message=error)

        return result['data']

    async def run_workflow_streaming(self, workflow_id, parameters=None, bot_id=None, app_id=None, ext=None):
        """ 执行工作流(流式)
        :param workflow_id:
        :param parameters:
        :param bot_id:
        :param app_id:
        :param ext:
        :return:
        """
        url = 'https://api.coze.cn/v1/workflow/stream_run'

        headers = await self._get_headers()
        params = {'workflow_id': workflow_id}
        if parameters is not None:
            params['parameters'] = parameters
        if bot_id is not None:
            params['bot_id'] = bot_id
        if app_id is not None:
            params['app_id'] = app_id
        if ext is not None:
            params['ext'] = ext

        async with aiohttp.ClientSession() as session:
            async with session.post(url, headers=headers, json=params) as resp:
                while True:
                    data = await resp.content.read(1024)
                    if not data:
                        break

                    try:
                        yield data.decode('utf-8')
                    except UnicodeDecodeError:
                        continue


if __name__ == "__main__":

    result = asyncio.run(Coze().create_conversation())
    print(result)
